# Copyright (c) OpenMMLab. All rights reserved.
import sys
from lmdeploy.utils import get_max_batch_size
import os
import argparse
from lmdeploy.cli.utils import (ArgumentHelper)
from constants import VERSION


class LlmProxy(object):     
    _help = 'Serve LLMs with openai API or triton server.'
    _desc = _help
    parser = argparse.ArgumentParser(prog='llm_proxy',
                                     description=_desc,
                                     add_help=True)
    parser.add_argument('-v',
                        '--version',
                        action='version',
                        version=VERSION)

    parser.add_argument('--server-name', type=str, default='0.0.0.0', help='Host ip for proxy serving')
    parser.add_argument('--server-port', type=int, default=8000, help='Server port of the proxy')
    parser.add_argument('--strategy',
                        type=str,
                        choices=['random', 'min_expected_latency', 'min_observed_latency', 'min_input_token'],
                        default='min_input_token',
                        help='the strategy to dispatch requests to nodes')
    ArgumentHelper.api_keys(parser)
    ArgumentHelper.ssl(parser)
    ArgumentHelper.log_level(parser)

def proxy_from_yaml(config_file: str=None):
    if config_file is None:
        current_dir = os.path.dirname(os.path.abspath(__file__))
        config_file = os.path.join(current_dir, 'config', 'proxy.yaml')

    if not os.path.exists(config_file):
        print(f'ERROR:  config file: {config_file} not found!')
        return        
    import yaml
    with open(config_file, 'r') as f:
        config = yaml.safe_load(f)
        from proxy_server import proxy
        proxy(**config)
def proxy():
    """LLM proxy with restful api using fastapi."""
    args = sys.argv[1:]
    parser = LlmProxy.parser  
    args = parser.parse_args()
    print(args)
    from proxy_server import proxy as api_proxy
    api_proxy(server_name=args.server_name,
          server_port=args.server_port,
          strategy=args.strategy,
          api_keys=args.api_keys,
          ssl=args.ssl,
          log_level=args.log_level.upper())    

if __name__ == '__main__':
    proxy()
'''    
    import fire
    fire.Fire(proxy_from_yaml)      
'''